# app/yolo_demo.py
from ultralytics import YOLO
from PIL import Image, ImageDraw
from pathlib import Path

class Detector:
    def __init__(self, weights: str = "yolov8n.pt", device: str | None = None):
        self.model = YOLO(weights)      # 首次会自动下载
        self.device = device            # None=自动（优先CUDA）

    def predict(self, image_path: str, conf: float = 0.25, iou: float = 0.45, imgsz: int = 1280):
        results = self.model.predict(
            source=str(image_path), conf=conf, iou=iou, imgsz=imgsz,
            device=self.device, verbose=False
        )[0]
        dets = []
        for b in results.boxes:
            xyxy = b.xyxy[0].tolist()
            score = float(b.conf[0])
            cls_id = int(b.cls[0])
            cls_name = self.model.names.get(cls_id, str(cls_id))
            dets.append({"bbox_xyxy": xyxy, "score": score, "cls_id": cls_id, "cls_name": cls_name})
        return dets

def draw_boxes(pil_img: Image.Image, dets: list[dict]) -> Image.Image:
    img = pil_img.copy()
    draw = ImageDraw.Draw(img)
    for d in dets:
        x1, y1, x2, y2 = d["bbox_xyxy"]
        draw.rectangle([x1, y1, x2, y2], outline=(255, 0, 0), width=3)
        draw.text((x1+3, y1+3), f'{d["cls_name"]} {d["score"]:.2f}')
    return img

# app/yolo_demo.py 追加
from functools import lru_cache
import numpy as np
import cv2
from PIL import Image

@lru_cache(maxsize=1)
def _get_model(weights: str = "yolov8n.pt"):
    from ultralytics import YOLO
    return YOLO(weights)

def run_yolo_on_image(pil_img: Image.Image, conf: float = 0.25, iou: float = 0.45,
                      imgsz: int = 1280, weights: str = "yolov8n.pt", device: str | None = None):
    model = _get_model(weights)
    res = model.predict(source=pil_img, conf=conf, iou=iou, imgsz=imgsz,
                        device=device, verbose=False)[0]
    dets = []
    for b in res.boxes:
        x1, y1, x2, y2 = [float(v) for v in b.xyxy[0].tolist()]
        name = model.names.get(int(b.cls[0]), str(int(b.cls[0])))
        dets.append({"name": name, "conf": float(b.conf[0]), "box": [x1, y1, x2, y2]})
    plotted_bgr = res.plot()                      # numpy BGR
    plotted_rgb = cv2.cvtColor(plotted_bgr, cv2.COLOR_BGR2RGB)
    return plotted_rgb, dets

# 让 draw_boxes 同时兼容 'bbox_xyxy' 和 'box'
def draw_boxes(pil_img: Image.Image, dets: list[dict]) -> Image.Image:
    img = pil_img.copy()
    from PIL import ImageDraw
    draw = ImageDraw.Draw(img)
    for d in dets:
        coords = d.get("bbox_xyxy") or d.get("box")
        x1, y1, x2, y2 = coords
        draw.rectangle([x1, y1, x2, y2], outline=(255, 0, 0), width=3)
        label = (d.get("cls_name") or d.get("name") or "obj") + f' {d.get("score", d.get("conf", 0.0)):.2f}'
        draw.text((x1 + 3, y1 + 3), label)
    return img


# --- add to app/yolo_demo.py ---
from functools import lru_cache
import cv2
from PIL import Image

@lru_cache(maxsize=1)
def _get_model(weights: str = "yolov8n.pt"):
    from ultralytics import YOLO
    return YOLO(weights)

def run_yolo_on_image(pil_img: Image.Image, conf: float = 0.25, iou: float = 0.45,
                      imgsz: int = 1280, weights: str = "yolov8n.pt", device: str | None = None):
    model = _get_model(weights)
    res = model.predict(source=pil_img, conf=conf, iou=iou, imgsz=imgsz,
                        device=device, verbose=False)[0]
    dets = []
    for b in res.boxes:
        x1, y1, x2, y2 = [float(v) for v in b.xyxy[0].tolist()]
        name = model.names.get(int(b.cls[0]), str(int(b.cls[0])))
        dets.append({"name": name, "conf": float(b.conf[0]), "box": [x1, y1, x2, y2]})
    plotted_bgr = res.plot()
    import numpy as np
    plotted_rgb = cv2.cvtColor(plotted_bgr, cv2.COLOR_BGR2RGB)
    return plotted_rgb, dets
